from flask import Flask, request, jsonify, render_template_string
import os
import csv
import torch
from transformers import T5Tokenizer, T5Model
import webbrowser
import threading

app = Flask(__name__)

# 配置 T5 模型
model_name = "t5-base"
device = "cuda" if torch.cuda.is_available() else "cpu"
tokenizer = T5Tokenizer.from_pretrained(model_name, legacy=False)
model = T5Model.from_pretrained(model_name).to(device)

# 定义函数：读取文件内容
def read_code_from_file(file_path):
    with open(file_path, "r", encoding="utf-8") as f:
        code = f.read().strip()
    return code

# 定义函数：编码代码并保存到 CSV
def encode_code_and_save_to_csv(code, csv_path):
    # 使用 T5 对代码进行编码
    inputs = tokenizer(code, return_tensors="pt", max_length=512, truncation=True).to(device)
    outputs = model.encoder(**inputs)
    embedding = outputs.last_hidden_state.mean(dim=1).squeeze().detach().cpu().tolist()

    # 将编码保存到 CSV 文件
    with open(csv_path, "w", newline="", encoding="utf-8") as csvfile:
        csv_writer = csv.writer(csvfile)
        csv_writer.writerow(["function_code", "vector"])
        csv_writer.writerow([code, embedding])

    return embedding
